import torch
from diffusers import StableDiffusion3InpaintPipeline
from PIL import  Image
import os
import numpy as np
# pipe = StableDiffusion3Pipeline.from_pretrained(
#     "/mnt/afs/luotianhang/models--stabilityai--stable-diffusion-3.5-medium/snapshots/b940f670f0eda2d07fbb75229e779da1ad11eb80", torch_dtype=torch.float16
# ).to("cuda")

# image = pipe(
#     "A shark holding a sign that says fucking world",
#     negative_prompt="",
#     num_inference_steps=28,
#     guidance_scale=7.0,
# ).images[0]
# image.save('./temp_sd3.png')

scale_index= 1

def generate_mask(x1y1x2y2,width,height):
        x1 = x1y1x2y2[0]
        y1 = x1y1x2y2[1]
        x2 = x1y1x2y2[2]
        y2 = x1y1x2y2[3]
        mask = np.zeros((368*scale_index,640*scale_index),dtype=np.uint8)
        mask[int(y1):int(y2),int(x1):int(x2)] = 255
        return Image.fromarray(mask).convert('L')

command = f'cp /mnt/afs/luotianhang/diffusers/examples/inpainting/sd_inpainting/checkpoint-29000/model.safetensors /mnt/afs/luotianhang/models--stabilityai--stable-diffusion-3.5-medium/snapshots/weight/transformer/diffusion_pytorch_model.safetensors'
os.system(command)
pipe = StableDiffusion3InpaintPipeline.from_pretrained(
    "/mnt/afs/luotianhang/models--stabilityai--stable-diffusion-3.5-medium/snapshots/weight", torch_dtype=torch.float16
)
pipe.safety_checker=None
pipe.to("cuda")

image = Image.open('/mnt/afs/luotianhang/bg.png').convert('RGB')
mem = {}
pad= 0
offset=300
mem['squirrel  ']=[792+offset, 307, 1135+offset, 750]
mem['rabbit  ']=[792+offset, 307, 1135+offset, 750]
mem['hedgehog  ']=[792+offset, 307, 1135+offset, 750]
mem['a furry dog']=[792+offset, 307, 1135+offset, 750]
mem['dog']=[792+offset, 307, 1135+offset, 750]
mem['a black tablet']=[752+offset, 718, 895+offset, 926]
mem['a laptop with an apple on the screen']=[792+offset, 307, 1135+offset, 750]
mem['computer']=[657+offset, 302, 1080+offset, 754] # 这个词肯能被理解成了集群
mem['cell phone ']=[775+offset, 430, 883+offset, 582]
mem['purple wallet ']=[775+offset, 430, 883+offset, 582]
mem['aeroplane']=[792-pad+offset, 307-pad, 1135+pad+offset, 750+pad]
        
width = image.width
height = image.height


count = 0
print('image2image')
for prompt , loc_x1y1x2y2 in mem.items():
    pad = 0
    x1=loc_x1y1x2y2[0]/width*(640*scale_index)
    y1=loc_x1y1x2y2[1]/height*(368*scale_index)
    x2=loc_x1y1x2y2[2]/width*(640*scale_index)
    y2=loc_x1y1x2y2[3]/height*(368*scale_index)

    x1 = max(x1-pad,0)
    y1 = max(y1-pad,0)
    x2 = min(x2+pad,640*scale_index)
    y2 = min(y2+pad,368*scale_index)

    loc_x1y1x2y2=[int(x1),int(y1),int(x2),int(y2)]
    mask_image = generate_mask(loc_x1y1x2y2,width,height)
        
    image = image.resize((640*scale_index,368*scale_index))
    res_image = pipe(
        prompt='realistic,best quality,  finely detailed, photorealistic, a whole body,completed,'+prompt,
        image=image, 
        mask_image=mask_image,
        height=368*scale_index,
        width=640*scale_index,
        num_inference_steps=28,
        guidance_scale=7.0,
        negative_prompt ="worst quality, low quality, normal quality, bad quality, blurry , ugly, chaos, 2D, cartoon"+'cartoon, illustration, 3d, sepia, painting, cartoons, sketch, (worst quality:2), ((monochrome)), ((grayscale:1.2)), (backlight:1.2), analog, analogphoto,freak,anomaly',
        ).images[0]
    compos = Image.blend(res_image.convert('RGBA'), mask_image.convert('RGBA'), alpha=0.5)
    save_dir=os.path.join('/mnt/afs/luotianhang/diffusers/examples/inpainting','demo','sd3','image2image')
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    res_image.save(f"{save_dir}/{str(count)+prompt.replace(' ','_')}.png")
    count+=1

print('text2image')

count= 0
for prompt  in mem.keys():

    prompt = prompt
    mask_image = generate_mask([0,0,image.width,image.height],image.width,image.height)
        
    
    image = image.resize((640*scale_index,368*scale_index))
    
    res_image = pipe(
        prompt='realistic,best quality,  finely detailed, '+prompt ,
        image=image, 
        mask_image=mask_image,
        height=368*scale_index,
        width=640*scale_index,
        num_inference_steps=28,
        guidance_scale=7.0,
        negative_prompt ="worst quality, low quality, normal quality, bad quality, blurry , ugly, chaos, 2D, cartoon",
        ).images[0]
    
    res_image=res_image.resize((640*scale_index,368*scale_index))
    
    save_dir=os.path.join('/mnt/afs/luotianhang/diffusers/examples/inpainting','demo','sd3','text2image')
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    res_image.save(f"{save_dir}/{str(count)+prompt}.png")
    count +=1